import numpy as np
import serial
import struct
from simple_hdlc import HDLC
from utils.point_mapping import point_map

# 电极驱动控制逻辑
class ElectrodeCtrl(object):

    def __init__(self, ipaddr_or_comport: str, port_or_baud: int, conn_type: str = 'net'):
        super(ElectrodeCtrl, self).__init__()
        # 根据连接类型初始化串口或网络连接
        if conn_type == 'net':
            self.s = serial.serial_for_url(f'socket://{ipaddr_or_comport}:{port_or_baud}', timeout=1)
        elif conn_type == 'com':
            self.s = serial.Serial(ipaddr_or_comport, port_or_baud, timeout=0.1)
        else:
            raise ValueError('Invalid conn_type')  # 抛出无效连接类型的错误
        self.h = HDLC(self.s)  # 初始化HDLC协议

    def _compress_array(self, voltage_array):
        # 压缩电压数组，将非零电压值转化为字节数组
        byte_array = np.packbits(voltage_array != 0)
        return byte_array.tobytes()  # 转换为字节类型

    def set_tick(self, pattern, pw: int, pp: int, pn: int, bp: int, volt: int='0', pulse_type: str='pos'):
        """
        打包Tick类的数据。

        参数:
        pattern : list - 8x8图案数组，每个元素是一个uint8
        pw : int    - 脉冲宽度，单位微秒
        pp : int    - 脉冲周期，单位微秒
        pn : int       - 脉冲数量
        bp : int    - 爆发周期，单位微秒
        volt_mv : int  - 输出电压，单位毫伏
        pulse_type : - 刺激类型，'pos'或'neg'

        返回:
        bytes - 打包后的数据
        """        
        if isinstance(pattern, list):
            pattern = np.array(pattern)  # 如果输入是列表，转换为numpy数组
        if pattern.shape != (8, 8):
            raise ValueError("pattern形状不正确,必须为8x8!")  # 检查数组大小是否为8x8

        # 创建格式字符串
        # 按照优化后的成员排序，首先是所有的uint32_t，然后是pattern数组和一个uint8_t
        fmt = f'<5I'
        
        # 将pattern从列表转换为bytes
        pattern_bytes = self._compress_array(point_map(pattern))

        # 打包数据
        packed_data = struct.pack(fmt, int(pw), int(pp), int(pn), int(bp), int(volt))
        packed_data += pattern_bytes 
        packed_data += bytes(b'\x00' if pulse_type =='pos' else b'\x01')

        #print(packed_data)
        self.h.sendFrame(packed_data)
        return packed_data

    def set_high_z(self, high_z_array):
        # 设置高阻态数组
        compressed_bytes = self._compress_array(high_z_array)  # 压缩高阻态数组
        self.h.sendFrame(compressed_bytes)  # 发送压缩后的字节数据

    def on(self):
        # 打开电极，设置所有电极为低阻态
        self.set_high_z(np.ones((8, ), dtype=np.uint8))

    def off(self):
        # 关闭电极，设置所有电极为高阻态
        self.set_high_z(np.zeros((8, ), dtype=np.uint8))
